package com.witmore.neutron.ai.cognitive.service;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.witmore.neutron.ai.api.dto.MessageDTO;
import com.witmore.neutron.ai.db.entity.AgentDO;
import com.witmore.neutron.ai.db.entity.AgentSessionRecordDO;
import com.witmore.neutron.ai.db.entity.AgentSessionRecordDetailDO;
import com.witmore.neutron.ai.db.mapper.AgentMapper;
import com.witmore.neutron.ai.db.mapper.AgentSessionRecordDetailMapper;
import com.witmore.neutron.ai.db.mapper.AgentSessionRecordMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.content.Media;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.CollectionUtils;
import org.springframework.util.ObjectUtils;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

/**
 * Description: AgentService
 * Author: 豆子高
 * Date: 2025/6/5 17:01
 */
@Slf4j
@Service
public class AgentService {

    @Autowired
    private AgentMapper agentMapper;
    @Autowired
    private AgentSessionRecordMapper agentSessionRecordMapper;
    @Autowired
    private AgentSessionRecordDetailMapper agentSessionRecordDetailMapper;
    @Autowired
    private RedisTemplate<String, String> redisTemplate;

    @Transactional(rollbackFor = Exception.class)
    public void saveAll(String sessionId, List<Message> messages) {
        AgentSessionRecordDO agentSessionRecordDO = agentSessionRecordMapper.selectOne(new LambdaQueryWrapper<AgentSessionRecordDO>().eq(AgentSessionRecordDO::getSessionId, sessionId));
        if (agentSessionRecordDO == null) {
            log.error("Session record not found for sessionId: {}", sessionId);
            return;
        }
        messages.forEach(message -> {
            // 需要过滤一下messages消息 去除库内已有数据 根据元数据值 key="db_id"来去重
            Map<String, Object> metadata = message.getMetadata();
            if (metadata != null && metadata.containsKey("db_id")) {
                return;
            }
            AgentSessionRecordDetailDO agentSessionRecordDetailDO = new AgentSessionRecordDetailDO();
            agentSessionRecordDetailDO.setUserId(agentSessionRecordDO.getUserId());
            agentSessionRecordDetailDO.setAgentId(agentSessionRecordDO.getAgentId());
            agentSessionRecordDetailDO.setAgentCode(agentSessionRecordDO.getAgentCode());
            agentSessionRecordDetailDO.setAgentName(agentSessionRecordDO.getAgentName());
            agentSessionRecordDetailDO.setSessionRecordId(agentSessionRecordDO.getId());
            agentSessionRecordDetailDO.setSessionId(sessionId);
            agentSessionRecordDetailDO.setSource(agentSessionRecordDO.getSource());
            agentSessionRecordDetailDO.setIsDeleted(false);
            agentSessionRecordDetailDO.setUpdateTime(LocalDateTime.now());
            agentSessionRecordDetailDO.setCreateTime(LocalDateTime.now());
            assert metadata != null;
            agentSessionRecordDetailDO.setShowFlag(metadata.get("showFlag") == null ? 1 : (Integer) metadata.get("showFlag"));
            agentSessionRecordDetailDO.setFiles(metadata.get("file") == null ? null : (String) metadata.get("file"));
            if (message instanceof UserMessage userMessage) {
                agentSessionRecordDetailDO.setRole(MessageType.USER.getValue());
                agentSessionRecordDetailDO.setContent(userMessage.getText());
                List<Media> media = userMessage.getMedia();
                agentSessionRecordDetailDO.setImages(MediaJsonConverter.mediaListToImagesList(media));
            } else if (message instanceof AssistantMessage assistantMessage) {
                // 由于要写进Usage需要去下重 以后若有更好的写法随时改TODO
                Map<String, Object> metadataA = assistantMessage.getMetadata();
                Object rowMetadataIdO = metadataA.get("id");
                if (rowMetadataIdO != null) {
                    String rowMetadataId = rowMetadataIdO.toString();
                    List<AgentSessionRecordDetailDO> agentSessionRecordDetailDOS = agentSessionRecordDetailMapper.selectList(new LambdaQueryWrapper<AgentSessionRecordDetailDO>()
                            .eq(AgentSessionRecordDetailDO::getRowMetadataId, rowMetadataId));
                    if (!CollectionUtils.isEmpty(agentSessionRecordDetailDOS)) {
                        // 存在则更新
                        AgentSessionRecordDetailDO agentSession = agentSessionRecordDetailDOS.get(0);
                        agentSession.setRowMetadata(JSONObject.toJSONString(metadataA));
                        agentSession.setUpdateTime(LocalDateTime.now());
                        agentSessionRecordDetailMapper.updateById(agentSession);
                        return;
                    }
                }
                agentSessionRecordDetailDO.setRowMetadataId(rowMetadataIdO == null ? null : rowMetadataIdO.toString());
                agentSessionRecordDetailDO.setRole(MessageType.ASSISTANT.getValue());
                agentSessionRecordDetailDO.setContent(assistantMessage.getText());
                agentSessionRecordDetailDO.setRowMetadata(JSONObject.toJSONString(metadataA));
            } else if (message instanceof SystemMessage systemMessage) {
                agentSessionRecordDetailDO.setRole(MessageType.SYSTEM.getValue());
                agentSessionRecordDetailDO.setContent(systemMessage.getText());
            } else if (message instanceof ToolResponseMessage toolResponseMessage) {
                agentSessionRecordDetailDO.setRole(MessageType.TOOL.getValue());
                agentSessionRecordDetailDO.setContent(JSONArray.toJSONString(toolResponseMessage.getResponses()));
                agentSessionRecordDetailDO.setShowFlag(0);
            } else {
                log.warn("Unknown message type: {}", message.getClass().getName());
            }
            agentSessionRecordDetailMapper.insert(agentSessionRecordDetailDO);

            // 如果有rowMetadataId，将数据存入Redis缓存，过期时间为1天
            if (agentSessionRecordDetailDO.getRowMetadataId() != null) {
                String cacheKey = "session_record_detail:" + agentSessionRecordDetailDO.getRowMetadataId();
                try {
                    // 将AgentSessionRecordDetailDO对象转换为JSON字符串存储
                    String jsonValue = JSONObject.toJSONString(agentSessionRecordDetailDO);
                    redisTemplate.opsForValue().set(cacheKey, jsonValue, 1, TimeUnit.DAYS);
                    log.info("已将会话记录详情缓存到Redis，key: {}", cacheKey);
                } catch (Exception e) {
                    log.error("缓存会话记录详情到Redis失败，key: {}", cacheKey, e);
                }
            }
        });
    }


    public List<Message> findSessionRecordById(String sessionId) {
        // 获取窗口值
        AgentSessionRecordDO agentSessionRecordDO = agentSessionRecordMapper.selectOne(new LambdaQueryWrapper<AgentSessionRecordDO>().eq(AgentSessionRecordDO::getSessionId, sessionId));
        AgentDO agentDO = agentMapper.selectById(agentSessionRecordDO.getAgentId());
        int maxMessages = 10;
        if (!ObjectUtils.isEmpty(agentDO.getAiConfig())) {
            JSONObject aiConfig = JSONObject.parseObject(agentDO.getAiConfig(), JSONObject.class);
            String dialogsNum = aiConfig.getString("dialogsNum");
            maxMessages = !ObjectUtils.isEmpty(dialogsNum) ? Integer.parseInt(dialogsNum) : 10;
        }
        Page<AgentSessionRecordDetailDO> sessionRecordDetailDOPage = agentSessionRecordDetailMapper.selectPage(new Page<>(1, maxMessages), new LambdaQueryWrapper<AgentSessionRecordDetailDO>()
                .eq(AgentSessionRecordDetailDO::getSessionId, sessionId)
                .eq(AgentSessionRecordDetailDO::getIsDeleted, 0)
                .orderByDesc(AgentSessionRecordDetailDO::getCreateTime)

        );
        if (ObjectUtils.isEmpty(sessionRecordDetailDOPage)) {
            return new ArrayList<>();
        }
        List<AgentSessionRecordDetailDO> agentSessionRecordDetailDOS = sessionRecordDetailDOPage.getRecords();
        // 按照id正序排一下
        agentSessionRecordDetailDOS.sort(Comparator.comparing(AgentSessionRecordDetailDO::getId));
        return agentSessionRecordDetailDOS.stream().map(agentSessionRecordDetailDO -> {
            Map<String, Object> metadata = new HashMap<>();
            metadata.put("db_id", agentSessionRecordDetailDO.getId());
            if (agentSessionRecordDetailDO.getRole().equalsIgnoreCase("SYSTEM")) {
                return SystemMessage.builder()
                        .metadata(metadata)
                        .text(agentSessionRecordDetailDO.getContent())
                        .build();
            }
            if (agentSessionRecordDetailDO.getRole().equalsIgnoreCase("USER")) {
                return UserMessage.builder()
                        .text(agentSessionRecordDetailDO.getContent())
                        .metadata(metadata)
                        .build();
            }
            if (agentSessionRecordDetailDO.getRole().equalsIgnoreCase("ASSISTANT")) {
                return new AssistantMessage(agentSessionRecordDetailDO.getContent(), metadata);
            }
            log.warn("Unknown role: {}, content: {}", agentSessionRecordDetailDO.getRole(), agentSessionRecordDetailDO.getContent());
            return new UserMessage(agentSessionRecordDetailDO.getContent());
        }).collect(Collectors.toList());
    }

    @Async("mvcTaskExecutor")
    public void asyncSaveSessionRecord(MessageDTO messageBO, ChatMemory chatMemory, Flux<ChatResponse> stream) {
        log.info("Start async processing on thread: {}", Thread.currentThread().getName());
        AtomicReference<String> rowMetadata = new AtomicReference<>("");
        stream.mapNotNull(chatResponse -> {
                    rowMetadata.set(JSONObject.toJSONString(chatResponse.getMetadata()));
                    return chatResponse.getResult().getOutput().getText();
                }).collect(StringBuilder::new, StringBuilder::append) // 收集所有片段
                .subscribe(completeResponse -> {
                    // 最终保存完整响应
                    chatMemory.add(messageBO.getSessionId(), new MetadataRichAssistantMessage(completeResponse.toString(), Map.of("metadata", rowMetadata.get())));
                    log.info("异步保存完成: {}", messageBO.getSessionId());
                }, error -> {
                    log.error("异步保存出错: {}", error.getMessage());
                    chatMemory.add(messageBO.getSessionId(), new AssistantMessage("[ERROR] " + error.getMessage()));
                });
    }




}
